"""
Created by joe on 2019/12/13

voc的xml标注 转化为 yolo txt标注
"""
import os
import sys

from bs4 import BeautifulSoup


class VocToYolo(object):
    """voc的xml标注转化为yolo格式"""

    def __init__(self):
        """
        init
        """
        self.classes = []

    @staticmethod
    def round_str(num, place):
        """
        四舍五入保留指定位数小数, 不够的小数位用字符串填充
        :param num:
        :param place:
        :return:
        """
        num = round(num, place)

        num_str = str(num)
        index = num_str.rindex(".")
        decimal_str = num_str[index + 1:]

        for _ in range(6 - len(decimal_str)):
            num_str += "0"

        return num_str

    def transfer(self, label_path):
        """
        transfer
        :param label_path: 标签路径
        :return:
        """
        for filename in os.listdir(label_path):
            name_arr = filename.split(".")

            if not name_arr[1] == "xml":
                continue

            xml_path = "{0}/{1}".format(label_path, filename)

            with open(xml_path, "r") as fi:
                document = BeautifulSoup(fi.read(), "html.parser")
                width = int(document.select("size width")[0].text)
                height = int(document.select("size height")[0].text)

                for obj_el in document.select("annotation object"):
                    class_name = obj_el.select("name")[0].text

                    if class_name not in self.classes:
                        self.classes.append(class_name)

                    x_min = int(obj_el.select("bndbox xmin")[0].text)
                    y_min = int(obj_el.select("bndbox ymin")[0].text)
                    x_max = int(obj_el.select("bndbox xmax")[0].text)
                    y_max = int(obj_el.select("bndbox ymax")[0].text)

                    line = "{0} {1} {2} {3} {4}\n".format(
                        self.classes.index(class_name),
                        self.round_str(((x_max - x_min) / 2 + x_min) / width, 6),
                        self.round_str(((y_max - y_min) / 2 + y_min) / height, 6),
                        self.round_str((x_max - x_min) / width, 6),
                        self.round_str((y_max - y_min) / height, 6)
                    )

                    with open("{0}/{1}.txt".format(label_path, name_arr[0]), "a") as fi2:
                        fi2.write(line)
            # os.remove(xml_path)  # 这里最好不要删除，如有需要手动删除，留个后悔药

        with open("{0}/classes.txt".format(label_path), "w") as fi:
            for class_name in self.classes:
                fi.write(class_name + "\n")


if __name__ == "__main__":
    """main"""
    obj = VocToYolo()
    obj.transfer(sys.argv[1])
